import cv2
import numpy as np
from time import *
import acl
from aclnet import aclnet
from aclnet import acldvpp
from aclnet.constant import check_ret, AscendContext

classes = 59
img_width = 2048
img_height = 1536
device_id = 0
decoder = 'opencv' # use --insert_op_conf=aipp_rgb.cfg
# decoder = 'dvpp' # use --insert_op_conf=aipp_yuv.cfg

def PostProcessing(result):
    color_map = np.zeros((classes, 3), dtype=np.uint8)
    for i in range(classes):
        color_map[i][0] = np.random.randint(0, 255, dtype=np.uint8)
        color_map[i][1] = np.random.randint(0, 255, dtype=np.uint8)
        color_map[i][2] = np.random.randint(0, 255, dtype=np.uint8)

    img_result = np.zeros((img_height//4, img_width//4, 3), dtype=np.uint8)

    for h in range(result.shape[1]):
        for w in range(result.shape[2]):
            color_index = np.argmax(result[:, h, w])
            img_result[h, w, :] = color_map[color_index]
    cv2.imwrite('result.jpg', img_result)

if __name__ == "__main__":
    with AscendContext(device_id) as context:
        # load model
        model = aclnet.Net(context, './model/OCRNet.om')
        image = []
        if decoder == 'opencv':
            image = cv2.imread("data/moto.jpg")
            image = cv2.resize(image, (img_width, img_height))
        elif decoder == 'dvpp':
            dvpp = acldvpp.Dvpp(context)
            image = dvpp.imread("data/moto.jpg")  # just support jpg
            image = dvpp.resize(image, img_width, img_height)

        begin_time = time()
        result = model.run([image], decoder, [(1, classes, img_height // 4, img_width // 4)])
        PostProcessing(result[0][0])
        end_time = time()
        print('total run time:', end_time - begin_time)
        del(image)
        del(model)
        del(dvpp)